set pipeline_dtype default value to params_dtype in megatron_eagle TransformerConfig #291
Conversation
Signed-off-by: Ye Yu <[email protected]>
WalkthroughUpdated dict_to_config in the Megatron-Eagle speculative plugin to set TransformerConfig.pipeline_dtype using architecture_config["torch_dtype"] via getattr(torch, ...). Previously it was None. No function signature or other field assignments changed. No public API modifications. Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Poem
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
modelopt/torch/speculative/plugins/megatron_eagle.py (1)
78-96: Deduplicate dtype resolution and accept torch.dtype directly.Resolve
torch_dtypeonce and reuse to avoid drift if resolution logic changes. Also handle the case wherearchitecture_config["torch_dtype"]is already atorch.dtype.Apply:
@@ def dict_to_config( @@ ): """Helper function to convert a dictionary to TransformerConfig.""" - config = TransformerConfig( + dtype = architecture_config.get("torch_dtype") + if isinstance(dtype, str): + dtype = getattr(torch, dtype) + elif not isinstance(dtype, torch.dtype): + raise ValueError(f"Unsupported torch_dtype: {dtype!r}") + + config = TransformerConfig( @@ - params_dtype=getattr(torch, architecture_config["torch_dtype"]), - pipeline_dtype=getattr(torch, architecture_config["torch_dtype"]), + params_dtype=dtype, + pipeline_dtype=dtype, num_layers=architecture_config.get("num_hidden_layers"),
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (1)
modelopt/torch/speculative/plugins/megatron_eagle.py(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: linux
- GitHub Check: wait-checks / wait
- GitHub Check: code-quality
- GitHub Check: build-docs
🔇 Additional comments (2)
modelopt/torch/speculative/plugins/megatron_eagle.py (2)
92-94: Good fix: pipeline_dtype now aligned with params_dtype.Setting
pipeline_dtypestops PP from being blocked byNoneand matchesparams_dtypeby construction. Looks correct.
79-108: Verify pipeline_dtype support in TransformerConfig
Ensure themegatron.coreversion in your CI environment accepts thepipeline_dtypeparameter inTransformerConfig.__init__; if not, remove or guard its use.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #291 +/- ##
==========================================
- Coverage 73.95% 73.94% -0.01%
==========================================
Files 172 172
Lines 17405 17405
==========================================
- Hits 12871 12870 -1
- Misses 4534 4535 +1 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
What does this PR do?
Type of change: Bug fix
Overview:
pipeline_dtype was set to None in current main which blocks PP usage in Megatron. We set to params_dtype to fix it.
Usage
# Add a code snippet demonstrating how to use thisTesting
Before your PR is "Ready for review"
Additional Information
Summary by CodeRabbit